from django.shortcuts import render, redirect, get_object_or_404
from django.http import HttpResponse, JsonResponse, StreamingHttpResponse
from django.views.decorators.csrf import csrf_exempt
from django.urls import reverse
from django.utils.text import slugify
from .KnowledgeClinet import KnowledgeClient, CollectionNotFoundError, DocumentNotFoundError
from .RerankClient import RerankClient
from .EmbeddingClient import EmbeddingClient
import os
import json
from dotenv import load_dotenv
import re

from .llm import LLM

load_dotenv()

# 初始化知识库客户端
knowledge_client = KnowledgeClient(
    persist_directory=os.path.join(os.path.dirname(__file__), "db"),
    embedding_client=EmbeddingClient(api_key=os.getenv("OPENAI_EMBEDDING_API_KEY"),
                                     base_url=os.getenv("OPENAI_EMBEDDING_API_URL"),
                                     model=os.getenv("OPENAI_EMBEDDING_MODEL_NAME")
                                     ),
    rerank_client=RerankClient(api_key=os.getenv("OPENAI_RARANKER_API_KEY"),
                               base_url=os.getenv("OPENAI_RARANKER_API_URL"),
                               model=os.getenv("OPENAI_RARANKER_MODEL_NAME"))
)


# 知识库管理首页 - 显示所有集合
def knowledge_index(request):
    collections = knowledge_client.list_collections()
    return render(request, 'knowledge/index.html', {
        'collections': collections,
        'active_page': 'index'
    })


# 创建新集合
def create_collection(request):
    if request.method == 'POST':
        collection_name = request.POST.get('name')
        if collection_name and collection_name != '':
            # 获取并解析元数据
            metadata_json = request.POST.get('metadata', '{}')
            try:
                metadata = json.loads(metadata_json) if metadata_json.strip() else {}
            except json.JSONDecodeError:
                return render(request, 'knowledge/create_collection.html', {
                    'error': '元数据格式不正确，应为有效的JSON',
                    'active_page': 'create_collection'
                })

            # 验证元数据是否为字典
            if not isinstance(metadata, dict):
                return render(request, 'knowledge/create_collection.html', {
                    'error': '元数据必须是JSON对象（键值对）',
                    'active_page': 'create_collection'
                })

            # 验证集合名称格式
            if not re.match(r'^[a-zA-Z0-9][a-zA-Z0-9._-]{1,510}[a-zA-Z0-9]$', collection_name):
                return render(request, 'knowledge/create_collection.html', {
                    'error': '集合名称只能包含字母、数字、点、下划线和连字符，长度3-512字符，且必须以字母或数字开头和结尾',
                    'active_page': 'create_collection'
                })
            try:
                knowledge_client.create_collection(collection_name, get_or_create=False, metadata=metadata)
                return redirect('knowledge:index')
            except ValueError:
                return render(request, 'knowledge/create_collection.html', {
                    'error': '集合名称已存在',
                    'active_page': 'create_collection'
                })
        else:
            return render(request, 'knowledge/create_collection.html', {
                'error': '集合名称不能为空',
                'active_page': 'create_collection'
            })
    return render(request, 'knowledge/create_collection.html', {
        'active_page': 'create_collection'
    })


# 删除集合
def delete_collection(request, collection_name):
    if request.method == 'POST':
        try:
            knowledge_client.delete_collection(collection_name)
            return redirect('knowledge:index')
        except CollectionNotFoundError:
            return HttpResponse('集合不存在', status=404)
    return render(request, 'knowledge/confirm_delete.html', {
        'item_name': collection_name,
        'item_type': '集合',
        'cancel_url': reverse('knowledge:index'),
        'active_page': 'index'
    })


# 集合详情 - 显示集合中的文档
def collection_detail(request, collection_name):
    try:
        # 获取集合中的所有文档
        collection = knowledge_client.get_collection(collection_name)
        documents = collection.get()
        # 格式化文档数据
        formatted_docs = []
        for doc_id, text, metadata in zip(
                documents['ids'], documents['documents'], documents['metadatas']
        ):
            formatted_docs.append({
                'id': doc_id,
                'title': metadata.get('title', '无标题'),
                'text': text[:100] + '...' if len(text) > 100 else text,
                'metadata': metadata,
                'full_text': text
            })
        return render(request, 'knowledge/collection_detail.html', {
            'collection_name': collection_name,
            'documents': formatted_docs,
            'active_page': 'index'
        })
    except CollectionNotFoundError:
        return HttpResponse('集合不存在', status=404)


# 添加文档到集合
def add_document(request, collection_name):
    if request.method == 'POST':
        title = request.POST.get('title')
        content = request.POST.get('content')
        metadata_json = request.POST.get('metadata', '{}')

        try:
            metadata = json.loads(metadata_json)
        except json.JSONDecodeError:
            return render(request, 'knowledge/add_document.html', {
                'collection_name': collection_name,
                'error': '元数据格式不正确，应为JSON',
                'active_page': 'index'
            })

        # 添加标题到元数据
        metadata['title'] = title

        try:
            # 添加文档到知识库
            doc_id = knowledge_client.add_documents(
                collection_name=collection_name,
                texts=[content],
                metadatas=[metadata]
            )[0]
            return redirect('knowledge:collection_detail', collection_name=collection_name)
        except CollectionNotFoundError:
            return HttpResponse('集合不存在', status=404)
        except Exception as e:
            return render(request, 'knowledge/add_document.html', {
                'collection_name': collection_name,
                'error': f'添加文档失败: {str(e)}',
                'active_page': 'index'
            })

    return render(request, 'knowledge/add_document.html', {
        'collection_name': collection_name,
        'active_page': 'index'
    })


# 编辑文档
def edit_document(request, collection_name, document_id):
    try:
        # 获取文档当前内容
        text, metadata = knowledge_client.get_document(collection_name, document_id)
        title = metadata.get('title', '')

        if request.method == 'POST':
            new_title = request.POST.get('title')
            new_content = request.POST.get('content')
            new_metadata_json = request.POST.get('metadata', '{}')

            try:
                new_metadata = json.loads(new_metadata_json)
            except json.JSONDecodeError:
                return render(request, 'knowledge/edit_document.html', {
                    'collection_name': collection_name,
                    'document_id': document_id,
                    'title': new_title,
                    'content': new_content,
                    'metadata': new_metadata_json,
                    'error': '元数据格式不正确，应为JSON',
                    'active_page': 'index'
                })

            # 更新标题到元数据
            new_metadata['title'] = new_title

            try:
                knowledge_client.update_document(
                    collection_name=collection_name,
                    document_id=document_id,
                    text=new_content,
                    metadata=new_metadata
                )
                return redirect('knowledge:collection_detail', collection_name=collection_name)
            except Exception as e:
                return render(request, 'knowledge/edit_document.html', {
                    'collection_name': collection_name,
                    'document_id': document_id,
                    'title': new_title,
                    'content': new_content,
                    'metadata': new_metadata_json,
                    'error': f'更新文档失败: {str(e)}',
                    'active_page': 'index'
                })

        # 将元数据转换为JSON字符串以便在表单中显示
        metadata_json = json.dumps(metadata, ensure_ascii=False, indent=2)

        return render(request, 'knowledge/edit_document.html', {
            'collection_name': collection_name,
            'document_id': document_id,
            'title': title,
            'content': text,
            'metadata': metadata_json,
            'active_page': 'index'
        })
    except CollectionNotFoundError:
        return HttpResponse('集合不存在', status=404)
    except DocumentNotFoundError:
        return HttpResponse('文档不存在', status=404)


# 删除文档
def delete_document(request, collection_name, document_id):
    if request.method == 'POST':
        try:
            knowledge_client.delete_document(collection_name, document_id)
            return redirect('knowledge:collection_detail', collection_name=collection_name)
        except (CollectionNotFoundError, DocumentNotFoundError):
            return HttpResponse('集合或文档不存在', status=404)

    # 获取文档标题用于确认页面
    try:
        text, metadata = knowledge_client.get_document(collection_name, document_id)
        title = metadata.get('title', '无标题文档')
        return render(request, 'knowledge/confirm_delete.html', {
            'item_name': title,
            'item_type': '文档',
            'cancel_url': reverse('knowledge:collection_detail', args=[collection_name]),
            'active_page': 'index'
        })
    except (CollectionNotFoundError, DocumentNotFoundError):
        return HttpResponse('集合或文档不存在', status=404)


# 查看文档详情
def view_document(request, collection_name, document_id):
    try:
        text, metadata = knowledge_client.get_document(collection_name, document_id)
        title = metadata.get('title', '无标题文档')
        return render(request, 'knowledge/view_document.html', {
            'collection_name': collection_name,
            'document_id': document_id,
            'title': title,
            'content': text,
            'metadata': metadata,
            'active_page': 'index'
        })
    except CollectionNotFoundError:
        return HttpResponse('集合不存在', status=404)
    except DocumentNotFoundError:
        return HttpResponse('文档不存在', status=404)


# 搜索知识库
def search_knowledge(request, collection_name=None):
    query = request.GET.get('q', '')
    results = []

    if query:
        if collection_name:
            # 搜索特定集合
            try:
                results = knowledge_client.search(
                    collection_name=collection_name,
                    query=query,
                    n_results=20,
                    rerank=True
                )
            except CollectionNotFoundError:
                return HttpResponse('集合不存在', status=404)
        else:
            # 搜索所有集合
            collections = knowledge_client.list_collections()
            for coll in collections:
                try:
                    coll_results = knowledge_client.search(
                        collection_name=coll,
                        query=query,
                        n_results=5,
                        rerank=True
                    )
                    # 添加集合名称到结果中
                    for res in coll_results:
                        res.metadata['collection'] = coll
                    results.extend(coll_results)
                except Exception:
                    continue
            # 跨集合排序结果
            results.sort(key=lambda x: x.score, reverse=True)

    return render(request, 'knowledge/search_results.html', {
        'query': query,
        'results': results,
        'collection_name': collection_name,
        'active_page': 'search'
    })


# ai搜索知识库
def ai_search_knowledge(request):
    collection_name = request.GET.get('collection_name')
    query = request.GET.get('q', '')
    n_results = int(request.GET.get('n', 5))  # 获取返回数量，默认5
    min_score = float(request.GET.get('min_score', 0.0))  # 获取相似度阈值，默认0.0
    request_type = request.GET.get('type', 'all')  # 请求类型：all, kb_only, ai_only
    results = []
    answer = ""

    if query:
        try:
            # 如果请求类型是 all 或 kb_only，获取知识库结果
            if request_type in ['all', 'kb_only', 'ai_only']:
                if collection_name:
                    # 搜索特定集合
                    results = knowledge_client.search(
                        collection_name=collection_name,
                        query=query,
                        n_results=n_results,
                        rerank=True
                    )
                else:
                    # 搜索所有集合
                    collections = knowledge_client.list_collections()
                    for coll in collections:
                        try:
                            coll_results = knowledge_client.search(
                                collection_name=coll,
                                query=query,
                                n_results=n_results,
                                rerank=True
                            )
                            # 添加集合名称到结果中
                            for res in coll_results:
                                res.metadata['collection'] = coll
                            results.extend(coll_results)
                        except Exception:
                            continue
                    # 跨集合排序结果并限制数量
                    results.sort(key=lambda x: x.score, reverse=True)
                    results = results[:n_results]

                # 根据相似度阈值过滤结果
                results = [res for res in results if res.score >= min_score]
                # print(results)

            # 如果请求类型是 all 或 ai_only，获取AI回答
            if request_type in ['all', 'ai_only']:
                llm = LLM()
                # 如果是AJAX请求，返回流式响应
                if request.headers.get('X-Requested-With') == 'XMLHttpRequest' or request.GET.get('ajax') == 'true':
                    def stream_answer():
                        context = "\n".join([f"文档名：{res.metadata.get("title")}（相似度{res.score:.2f}）：{res.text}" for res in
                                             results]) if results else "无相关知识库内容"
                        messages = [
                            {"role": "system",
                             "content": """
                             角色定位：
你是一名高效的知识库助手，具备深度理解和快速检索能力，能够根据用户提供的上下文精准回答问题。

任务描述：
你的主要职责是依据存储在知识库中的文档资料，帮助用户找到所需的信息。请确保你的回答简洁、准确，并且紧密基于用户提问的相关上下文。

操作指南：
1. 仔细分析用户的提问，理解其核心需求。
2. 利用知识库资源，寻找最相关的答案片段。
4. 回答应控制在每点不超过500字，便于快速阅读。
5. 如果需要引用资料，请在文末单独列出所有参考过的文件索引或名称。
6. 避免提供不相关或未经验证的信息；所有信息必须有据可查。
7. 对于无法直接回答的问题，请引导用户提供更多细节或建议他们检查特定部分的文档。

格式要求：
- 使用标题（###）来区分不同部分。
- 代码或命令行指令需使用代码块（```）标记。
- 在文档末尾，使用独立的一行列出引用：“引用: #索引号, 文件名”。

限制条件：
仅针对与个人文档、Trae功能或本地知识库检索相关的问题提供服务。对于无关请求，请礼貌地告知用户此限制。
                             """},
                            {"role": "user", "content": f"问题：{query}\n上下文：{context}"}
                        ]
                        # print(messages)
                        # 使用流式响应
                        response = llm.chat(messages, temperature=0.3, max_tokens=500, stream=True)
                        for chunk in response:
                            if chunk.choices[0].delta.content:
                                # 直接返回文本内容，不使用SSE格式
                                yield chunk.choices[0].delta.content
                        # 不需要发送完成信号，流结束即表示完成

                    response = StreamingHttpResponse(stream_answer(), content_type='text/plain')
                    response['Cache-Control'] = 'no-cache'
                    return response
                else:
                    # 非AJAX请求，使用普通响应
                    context = "\n".join([f"文档{res.document_id}（相似度{res.score:.2f}）：{res.text}" for res in results])
                    messages = [
                        {"role": "system",
                         "content": """
                         你是一名高效的知识库助手，具备深度理解和快速检索能力，能够根据用户提供的上下文精准回答问题。

任务描述：
你的主要职责是依据存储在知识库中的文档资料，帮助用户找到所需的信息。请确保你的回答简洁、准确，并且紧密基于用户提问的相关上下文。

操作指南：
1. 仔细分析用户的提问，理解其核心需求。
2. 利用知识库资源，寻找最相关的答案片段。
4. 回答应控制在每点不超过500字，便于快速阅读。
5. 如果需要引用资料，请在文末单独列出所有参考过的文件索引或名称。
6. 避免提供不相关或未经验证的信息；所有信息必须有据可查。
7. 对于无法直接回答的问题，请引导用户提供更多细节或建议他们检查特定部分的文档。

格式要求：
- 使用标题（###）来区分不同部分。
- 代码或命令行指令需使用代码块（```）标记。
- 在文档末尾，使用独立的一行列出引用：“引用: #索引号, 文件名”。

限制条件：
仅针对与个人文档、Trae功能或本地知识库检索相关的问题提供服务。对于无关请求，请礼貌地告知用户此限制。

                         """},
                        {"role": "user", "content": f"问题：{query}\n上下文：{context}"}
                    ]
                    answer = llm.chat(messages, temperature=0.3, max_tokens=500)

        except CollectionNotFoundError:
            return HttpResponse('集合不存在', status=404)
        except Exception as e:
            return HttpResponse(f'搜索出错：{str(e)}', status=500)

    # 处理所有AJAX请求
    if request.headers.get('X-Requested-With') == 'XMLHttpRequest' or request.GET.get('ajax') == 'true':
        if request_type == 'kb_only':
            formatted_results = [{
                'document_id': res.metadata.get('title'),
                'text': res.text,
                'metadata': res.metadata,
                'score': res.score,
                'collection': res.metadata.get('collection', collection_name)
            } for res in results]
            return JsonResponse({'results': formatted_results})
        elif request_type == 'ai_only':
            if not results:
                return JsonResponse({'error': '没有找到相关知识库内容，无法生成AI回答'}, status=404)
            # AI流式响应已在前面处理
            return JsonResponse({})
        else:  # all类型
            formatted_results = [{
                'document_id': res.metadata.get('title'),
                'text': res.text,
                'metadata': res.metadata,
                'score': res.score,
                'collection': res.metadata.get('collection', collection_name)
            } for res in results]
            return JsonResponse({
                'results': formatted_results,
                'answer': answer
            })

    # 非AJAX请求才渲染页面
    all_collections = knowledge_client.list_collections()
    return render(request, 'knowledge/ai_search_results.html', {
        'query': query,
        'results': results,
        'answer': answer,
        'collection_name': collection_name,
        'collections': all_collections,
        'n_results': n_results,
        'min_score': min_score,
        'active_page': 'ai_search'
    })


# API接口 - 获取所有集合
@csrf_exempt
def api_collections(request):
    if request.method == 'GET':
        collections = knowledge_client.list_collections()
        return JsonResponse({'collections': collections})
    elif request.method == 'POST':
        data = json.loads(request.body)
        collection_name = data.get('name')
        if not collection_name:
            return JsonResponse({'error': '集合名称不能为空'}, status=400)
        try:
            knowledge_client.create_collection(collection_name, get_or_create=False)
            return JsonResponse({'status': 'success', 'collection': collection_name})
        except ValueError:
            return JsonResponse({'error': '集合已存在'}, status=400)


# API接口 - 操作集合中的文档
@csrf_exempt
def api_documents(request, collection_name):
    if request.method == 'GET':
        try:
            collection = knowledge_client.get_collection(collection_name)
            documents = collection.get()
            formatted_docs = []
            for doc_id, text, metadata in zip(
                    documents['ids'], documents['documents'], documents['metadatas']
            ):
                formatted_docs.append({
                    'id': doc_id,
                    'title': metadata.get('title', '无标题'),
                    'text': text,
                    'metadata': metadata
                })
            return JsonResponse({'documents': formatted_docs})
        except CollectionNotFoundError:
            return JsonResponse({'error': '集合不存在'}, status=404)

    elif request.method == 'POST':
        data = json.loads(request.body)
        content = data.get('content')
        metadata = data.get('metadata', {})

        if not content:
            return JsonResponse({'error': '文档内容不能为空'}, status=400)

        try:
            doc_ids = knowledge_client.add_documents(
                collection_name=collection_name,
                texts=[content],
                metadatas=[metadata]
            )
            return JsonResponse({'status': 'success', 'document_id': doc_ids[0]})
        except CollectionNotFoundError:
            return JsonResponse({'error': '集合不存在'}, status=404)


@csrf_exempt
def api_document_detail(request, collection_name, document_id):
    if request.method == 'GET':
        try:
            text, metadata = knowledge_client.get_document(collection_name, document_id)
            return JsonResponse({
                'id': document_id,
                'content': text,
                'metadata': metadata
            })
        except (CollectionNotFoundError, DocumentNotFoundError):
            return JsonResponse({'error': '集合或文档不存在'}, status=404)

    elif request.method == 'PUT':
        data = json.loads(request.body)
        content = data.get('content')
        metadata = data.get('metadata')

        if content is None:
            return JsonResponse({'error': '文档内容不能为空'}, status=400)

        try:
            knowledge_client.update_document(
                collection_name=collection_name,
                document_id=document_id,
                text=content,
                metadata=metadata
            )
            return JsonResponse({'status': 'success'})
        except (CollectionNotFoundError, DocumentNotFoundError):
            return JsonResponse({'error': '集合或文档不存在'}, status=404)

    elif request.method == 'DELETE':
        try:
            knowledge_client.delete_document(collection_name, document_id)
            return JsonResponse({'status': 'success'})
        except (CollectionNotFoundError, DocumentNotFoundError):
            return JsonResponse({'error': '集合或文档不存在'}, status=404)


# API接口 - 搜索知识库
@csrf_exempt
def api_search(request, collection_name=None):
    if request.method == 'GET':
        query = request.GET.get('q', '')
        if not query:
            return JsonResponse({'error': '查询参数q不能为空'}, status=400)

        n_results = int(request.GET.get('n', 10))
        rerank = request.GET.get('rerank', 'true').lower() == 'true'

        try:
            if collection_name:
                results = knowledge_client.search(
                    collection_name=collection_name,
                    query=query,
                    n_results=n_results,
                    rerank=rerank
                )
            else:
                collections = knowledge_client.list_collections()
                results = []
                for coll in collections:
                    coll_results = knowledge_client.search(
                        collection_name=coll,
                        query=query,
                        n_results=min(n_results, 5),
                        rerank=rerank
                    )
                    for res in coll_results:
                        res.metadata['collection'] = coll
                    results.extend(coll_results)
                results.sort(key=lambda x: x.score, reverse=True)
                results = results[:n_results]

            formatted_results = [{
                'document_id': res.document_id,
                'text': res.text,
                'metadata': res.metadata,
                'score': res.score
            } for res in results]

            return JsonResponse({'results': formatted_results})
        except CollectionNotFoundError:
            return JsonResponse({'error': '集合不存在'}, status=404)

    return JsonResponse({'error': '不支持的请求方法'}, status=405)
